/*
 * COPYRIGHT:	2010 Christopher Smith, licenced under the GNU GPL v2; see COPYING in the top level directory
 * PROJECT:		sdedit
 * FILE:		lsa.cpp
 * PURPOSE:		LSA (Local Security Authority) interface support
 */

namespace Lsa
{
using namespace Lookup;
using namespace WorkQueue;

void free_LsaHandle(HANDLE Handle)
{
	LsaClose(Handle);
}

auto_take<WCHAR> UNICODE_STRING2NewPWSTR(UNICODE_STRING *Us)
{
	SIZE_T Cch = Us->Length / sizeof(WCHAR) + 1;
	auto_szpod<WCHAR> Buff = auto_szpod_empty_ct(Cch + 1);
	RtlCopyMemory(Buff, Us->Buffer, Us->Length);
	Buff[Cch] = 0; //null terminate

	return Buff.release();
}

template <class T> struct LsaHeap
{
	typedef LsaHeap type;

	T *alloc(size_t length); //no alloc available
	void free(T *x)
	{
		LsaFreeMemory(x);
	}
};

class LsaPolicy : public auto_refctd_std_base
{
public:
	auto_free<VOID, free_LsaHandle> Handle;

private:
	LsaPolicy() {} //for special static constructors

	void OpenPolicy(UNICODE_STRING *SystemName, UCHAR Privileges)
	{
		OpText Ot(L"Opening a LSA policy", SystemName ? SystemName->Buffer : NULL);

		OBJECT_ATTRIBUTES Oa;
		RtlZeroLengthedStruct(Oa);

		RaiseOnNtError(LsaOpenPolicy(SystemName, &Oa, MAXIMUM_ALLOWED | StandardPrivilegedRights(Privileges), Handle.assign_indirect()));

		Ot;
	}

	void OpenPolicy(PCWSTR SystemName = NULL, UCHAR Privileges = 0)
	{		
		UNICODE_STRING US_SystemName;

		if(SystemName)
			RtlInitUnicodeString(&US_SystemName, const_cast<PWSTR>(SystemName));

		OpenPolicy(SystemName ? &US_SystemName : NULL, Privileges);
	}

public:
	LsaPolicy(HANDLE PreviousPolicy, UNICODE_STRING *DomainName, UCHAR Privileges)
	{
		RaiseOnNtError(LsaOpenTrustedDomainByName(PreviousPolicy, DomainName, MAXIMUM_ALLOWED | StandardPrivilegedRights(Privileges), Handle.assign_indirect()));
	}

	LsaPolicy(PCWSTR SystemName, UCHAR Privileges = 0)
	{
		OpenPolicy(SystemName, Privileges);
	}

	LsaPolicy(SDE_GENERIC_OPTS const &GenOpts, SDE_OPTS Opts, UCHAR Privileges)
	{
		SDE_OPT *Opt = Opts.Opts;

		PCWSTR SystemName = Opt[OiSystemName].Flags & SDE_OPT_SET ? Opt[OiSystemName].String : GenOpts.Host;

		if(!SystemName && GenOpts.Domain || Opt[OiDomain].Flags & SDE_OPT_SET)
		{
			PCWSTR Domain = Opt[OiDomain].Flags & SDE_OPT_SET ? Opt[OiDomain].String : GenOpts.Domain;

			auto_szpod<WCHAR> SystemName = Netapi::LookupDcHost(NULL, Domain);

			OpenPolicy(SystemName, Privileges);
		}
		else
			OpenPolicy(SystemName, Privileges);
	}

	static LsaPolicy *OpenNewPolicyByDomain(PCWSTR SearchHost, PCWSTR Domain)
	{
		auto_szpod<WCHAR> SystemName = Netapi::LookupDcHost(SearchHost, Domain);

		LsaPolicy *Ret = new LsaPolicy();

		Ret->OpenPolicy(SystemName, 0);

		return Ret;
	}

	static LsaPolicy *OpenNewPolicyByUCS(UNICODE_STRING *SystemName, UCHAR Privileges)
	{
		LsaPolicy *Ret = new LsaPolicy();

		Ret->OpenPolicy(SystemName, Privileges);
	}

	~LsaPolicy() {}
};

}//namespace Lsa

#include "stdsid.cpp"

namespace Lsa
{

struct LsaLookupSidWorkItem : LookupWorkItem<WCHAR>
{
	auto_szpod<SID> Sid;
	auto_refct<LsaPolicy> Policy;

	LsaLookupSidWorkItem(SID const *Sid, auto_refct<LsaPolicy> &Policy, CompletionTarget<WCHAR> *Ct) :
		Sid(Sid),
		Policy(Policy),
		LookupWorkItem(Ct)
	{}

	virtual void Work() override
	{
		OpText Ot(L"Looking up a SID via LSA");
		auto_ptr<LSA_REFERENCED_DOMAIN_LIST, LsaHeap> ReferencedDomain;
		auto_ptr<LSA_TRANSLATED_NAME, LsaHeap> TranslatedName;
		PSID Sids = Sid;

		RaiseOnNtError(LsaLookupSids(Policy->Handle, 1, &Sids, ReferencedDomain.assign_indirect(), TranslatedName.assign_indirect()));

		Result.realloc(ReferencedDomain->Domains->Name.Length + TranslatedName->Name.Length + 2 * sizeof(WCHAR));

		USHORT DomainCt = ReferencedDomain->Domains->Name.Length / sizeof(WCHAR);
		USHORT NameCt = TranslatedName->Name.Length / sizeof(WCHAR);

		//copy domain
		RtlCopyMemory(Result, ReferencedDomain->Domains->Name.Buffer, ReferencedDomain->Domains->Name.Length);
		//add backslash
		Result[DomainCt] = L'\\';
		//copy name
		RtlCopyMemory(Result + DomainCt + 1, TranslatedName->Name.Buffer, TranslatedName->Name.Length);
		//terminate
		Result[DomainCt + NameCt + 1] = L'\0';

		Ot;
	}
};

struct LsaLookupNameWorkItem : LookupWorkItem<SID>
{
	auto_szpod<WCHAR> Name;
	auto_refct<LsaPolicy> Policy;

	LsaLookupNameWorkItem(PCWSTR Name, auto_refct<LsaPolicy> &Policy, CompletionTarget<SID> *Ct) :
		Name(Name),
		Policy(Policy),
		LookupWorkItem(Ct)
	{}

	virtual void Work() override
	{
		OpText Ot(L"Looking up a name via LSA");
		auto_ptr<LSA_REFERENCED_DOMAIN_LIST, LsaHeap> ReferencedDomain;
		auto_ptr<LSA_TRANSLATED_SID, LsaHeap> Sid;
		UNICODE_STRING Name2;
		RtlInitUnicodeString(&Name2, Name);

		RaiseOnNtError(LsaLookupNames(Policy->Handle, 1, &Name2, ReferencedDomain.assign_indirect(), Sid.assign_indirect()));

		PSID ReferencedDomainSid = ReferencedDomain->Domains[0].Sid;

		Result = AddRidToSid(c_cast<SID *>(ReferencedDomainSid), Sid->RelativeId);

		Ot;
	}
};

struct SidTranslator : Lookup::Translator<SID>
{
	auto_refct<LsaPolicy> Policy;

	SidTranslator(LsaPolicy *Policy) :
		Policy(Policy)
	{}

	virtual void LookupData(CompletionTarget<WCHAR> *Ct, SID const *Sid) override
	{
		PostWorkItem(new LsaLookupSidWorkItem(Sid, Policy, Ct));
	}
	virtual void LookupName(CompletionTarget<SID> *Ct, PCWSTR Name) override
	{
		PostWorkItem(new LsaLookupNameWorkItem(Name, Policy, Ct));
	}
};

struct AccountEnumerator : Enumerator<SID>
{
	auto_refct<LsaPolicy> Policy;
	ULONG SidTypeFilter;

	AccountEnumerator(LsaPolicy *Policy, ULONG SidTypeFilter) :
		Policy(Policy),
		SidTypeFilter(SidTypeFilter)
	{}

	AccountEnumerator(DOMAIN_INFO const &Di, ULONG SidTypeFilter) :
		SidTypeFilter(SidTypeFilter)
	{
		if(Di.Flags & DOMAIN_INFO_HAS_HOST)
			Policy = new LsaPolicy(Di.Host);
		else
			Policy = LsaPolicy::OpenNewPolicyByDomain(Di.SearchHost, Di.DomainName);
	}

	virtual void WalkChildren(ChildrenWalkDest<SID> &Dest) override
	{} //lsa domains are flat

	virtual void WalkEntries(EntriesWalkDest<SID> &Dest) override
	{
		LSA_ENUMERATION_HANDLE Context = 0;
		auto_ptr<SID *, LsaHeap> EnumInfo;
		ULONG EnumInfoCt;
 
		NTSTATUS Status = STATUS_MORE_ENTRIES;

		while(Status == STATUS_MORE_ENTRIES)
		{
			Status = LsaEnumerateAccounts(Policy->Handle, &Context, EnumInfo.assign_indirect(), 0x1000, &EnumInfoCt);

			if(Status == STATUS_NO_MORE_ENTRIES)
				break;

			RaiseOnNtError(Status);

			//lookup names
			auto_ptr<LSA_REFERENCED_DOMAIN_LIST, LsaHeap> ReferencedDomains;
			auto_ptr<LSA_TRANSLATED_NAME, LsaHeap> TranslatedNames;
			PSID *Sids = (PSID *)EnumInfo.get();

			Status = LsaLookupSids(Policy->Handle, EnumInfoCt, Sids, ReferencedDomains.assign_indirect(), TranslatedNames.assign_indirect());

			RaiseOnNtError(Status); //note: wont't raise on STATUS_SOME_NOT_MAPPED; it's a SUCCESS code (no sense quitting if only some failed)

			for(ULONG i = 0; i < EnumInfoCt; i++)
				if(IsSidTypeInMask(TranslatedNames[i].Use, SidTypeFilter))
				{
					auto_szpod<WCHAR> Pwstr;
					if(TranslatedNames[i].Name.Length)
						Pwstr = UNICODE_STRING2NewPWSTR(&TranslatedNames[i].Name);
					else
						Pwstr = SidToText(EnumInfo[i]);

					Dest.EntryProc(Pwstr, *EnumInfo[i]);
				}
		}
	}
};

struct DomainAccountEnumerator : Lookup::DomainAccountEnumerator
{
	ULONG SidTypeFilter;

	DomainAccountEnumerator() :
		SidTypeFilter(SID_TYPE_USER | SID_TYPE_GROUP | SID_TYPE_ALIAS | SID_TYPE_COMPUTER)
	{}

	DomainAccountEnumerator(ULONG SidTypeFilter) :
		SidTypeFilter(SidTypeFilter)
	{}

	virtual Enumerator<SID> *NewEnum(DOMAIN_INFO const &Di) override
	{
		return new AccountEnumerator(Di, SidTypeFilter);
	}
};

struct DomainEnumerator : Lookup::DomainEnumerator
{
	auto_refct<LsaPolicy> Policy;

	DomainEnumerator(LsaPolicy *Policy) :
		Policy(Policy)
	{}

	virtual void WalkEntries(EntriesWalkDest<DOMAIN_INFO> &Dest) override
	{
		//get local acctdb info
		auto_ptr<POLICY_ACCOUNT_DOMAIN_INFO, LsaHeap> ADInfo;
		RaiseOnNtError(LsaQueryInformationPolicy(Policy->Handle, PolicyAccountDomainInformation, c_cast<PVOID *>(ADInfo.assign_indirect())));
		auto_szpod<WCHAR> ADName = UNICODE_STRING2NewPWSTR(&ADInfo->DomainName);

		//get primary domain info
		auto_ptr<POLICY_PRIMARY_DOMAIN_INFO, LsaHeap> PDInfo;
		RaiseOnNtError(LsaQueryInformationPolicy(Policy->Handle, PolicyPrimaryDomainInformation, c_cast<PVOID *>(PDInfo.assign_indirect())));

		auto_szpod<WCHAR> PriDomainName, DomainCtlrName;
		auto_refct<LsaPolicy> PriDomainPolicy;

		if(PDInfo->Sid)
		{
			PriDomainName = UNICODE_STRING2NewPWSTR(&PDInfo->Name);
			DomainCtlrName = Netapi::LookupDcHost(NULL, PriDomainName);
		}

		if(PDInfo->Sid && EqualSid(PDInfo->Sid, ADInfo->DomainSid)) //system IS a domain controller
		{
			PriDomainPolicy = Policy;

			DOMAIN_INFO PdDomainInfo = {DOMAIN_INFO_HAS_HOST | DOMAIN_INFO_NETDOMAIN,
				c_cast<SID *>(PDInfo->Sid), PriDomainName, DomainCtlrName};
			Dest.EntryProc(PriDomainName, PdDomainInfo);
		}
		else //system is not a domain controller
		{
			//enum local accounts
			auto_szpod<WCHAR> LocalDomainName = UNICODE_STRING2NewPWSTR(&ADInfo->DomainName);
			DOMAIN_INFO LocalDomainInfo = {DOMAIN_INFO_HAS_HOST | DOMAIN_INFO_LOCALDOMAIN,
				c_cast<SID *>(ADInfo->DomainSid), LocalDomainName, LocalDomainName};
			Dest.EntryProc(LocalDomainName, LocalDomainInfo);

			if(!PDInfo->Sid) //if machine doesn't belong to a domain, there's nothing else to do
				return;

			PriDomainPolicy = new LsaPolicy(DomainCtlrName, 0);

			//enum primary domain accounts
			DOMAIN_INFO PdDomainInfo = {DOMAIN_INFO_HAS_HOST | DOMAIN_INFO_NETDOMAIN,
				c_cast<SID *>(PDInfo->Sid), PriDomainName, DomainCtlrName};
			Dest.EntryProc(PriDomainName, PdDomainInfo);
		}
		
		LSA_ENUMERATION_HANDLE Context = 0;
		auto_ptr<LSA_TRUST_INFORMATION, LsaHeap> EnumInfo;
		ULONG EnumInfoCt;
 
		NTSTATUS Status = STATUS_MORE_ENTRIES;

		//go through each trusted domain
		while(Status == STATUS_MORE_ENTRIES)
		{
			Status = LsaEnumerateTrustedDomains(PriDomainPolicy->Handle, &Context, c_cast<PVOID *>(EnumInfo.assign_indirect()), 0x1000, &EnumInfoCt);

			if(Status == STATUS_NO_MORE_ENTRIES)
				break;

			RaiseOnNtError(Status);

			for(ULONG i = 0; i < EnumInfoCt; i++)
			{
				auto_szpod<WCHAR> TrustedDomainName = UNICODE_STRING2NewPWSTR(&EnumInfo[i].Name);

				DOMAIN_INFO TrustedDomainInfo = {DOMAIN_INFO_HAS_SEARCHHOST | DOMAIN_INFO_NETDOMAIN,
					c_cast<SID *>(EnumInfo->Sid), TrustedDomainName, NULL, DomainCtlrName};
				Dest.EntryProc(TrustedDomainName, TrustedDomainInfo);
			}
		}
	}

	virtual void WalkChildren(ChildrenWalkDest<DOMAIN_INFO> &Dest) override
	{}//LSA domains are flat: no children
};

}//namespace Lsa
